import os
import sys
import json
import base64
import argparse
from typing import List, Dict, Any
from openai import OpenAI

SYSTEM_PROMPT = """你是一个图文匹配专家。现在给你若干张图片，以及一个候选标题列表（仅限这些候选）。
你的任务：为每张图片从候选中选出**唯一**一个最合适的标题，形成一对一（双射）匹配：
- 每张图片恰好匹配 1 个标题；
- 每个标题最多被使用 1 次；
- 标题必须从候选列表中选择，不得自拟或改写；
- 如果实在没有合理匹配，请为该图片输出空字符串 ""（但仍保持所有图片各占一个结果，且不复用同一候选）。

匹配准则（从高到低）：
1) 视觉实体与布局高度一致（图表类型、坐标轴/图例、地图/表格/热力图/示意图等）。
2) 关键术语与可见文本一致（模型名、任务名、指标名、国家/城市名、维度标签等）。
3) 语义与粒度相符（越具体越好；若多候选都泛化，选最不容易冲突的那一条）。

输出格式**必须是合法 JSON**，不能包含任何额外文字或 Markdown。格式为：
{
  "assignments": [
    {"index": 0, "caption": "<从候选中挑选的一条或空字符串>"},
    {"index": 1, "caption": "..."}
    ...
  ]
}
其中 index 对应我提供的“图片顺序”（从 0 开始）。"""

USER_PROMPT_TEMPLATE = """请为下列图片与候选标题做一对一匹配：

候选标题（共 {k} 条，索引 0..{k_minus_1}）：
{candidates_block}

要求：
- 对于我按顺序提供的每张图片（index 0..{n_minus_1}），从候选中各选 1 条作为 caption；
- 同一候选不可被多张图片复用；
- 如无合适匹配，可为该图片输出空字符串 ""。

请只输出 JSON，格式：
{{
  "assignments": [
    {{"index": 0, "caption": "<候选或空串>"}},
    ...
  ]
}}
"""

def to_data_url_png(image_path: str) -> str:
    with open(image_path, "rb") as f:
        b64 = base64.b64encode(f.read()).decode("utf-8")
    return f"data:image/png;base64,{b64}"

def map_source_to_png_path(src_path: str, image_root: str) -> str:
    """
    将 source 路径映射到 image_root 下同名 .png：
    - xxx.pdf -> image_root/xxx.png
    - xxx.png -> 若 image_root/xxx.png 存在优先用之，否则用原路径
    - 其他后缀 -> 去后缀 + .png 到 image_root
    """
    base = os.path.basename(src_path)
    name, _ = os.path.splitext(base)
    candidate = os.path.join(image_root, f"{name}.png")
    if os.path.exists(candidate):
        return candidate
    if src_path.lower().endswith(".png") and os.path.exists(src_path):
        return src_path
    return candidate  # 可能不存在，上层决定处理

def strip_json_fence(s: str) -> str:
    s = s.strip()
    if s.startswith("```"):
        s = s.strip("`")
    first = s.find("{")
    last = s.rfind("}")
    if first != -1 and last != -1 and last >= first:
        return s[first:last+1]
    return s

def safe_json_loads(s: str):
    try:
        return json.loads(s)
    except Exception:
        return json.loads(strip_json_fence(s))

def build_messages(images_data_urls: List[str], candidates: List[str]) -> List[Dict[str, Any]]:
    content = []
    for url in images_data_urls:
        content.append({"type": "image_url", "image_url": {"url": url}})
    candidates_block = "\n".join([f"{i}. {c}" for i, c in enumerate(candidates)])
    text = USER_PROMPT_TEMPLATE.format(
        k=len(candidates),
        k_minus_1=max(len(candidates)-1, 0),
        candidates_block=candidates_block,
        n_minus_1=max(len(images_data_urls)-1, 0)
    )
    content.append({"type": "text", "text": text})
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": content}
    ]

def call_qvq(client: OpenAI, model: str, messages: list) -> str:
    """调用 qvq 模型接口（禁用流式处理）"""
    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        stream=False,  # 将流式处理禁用
    )
    return completion.choices[0].message.content


def main():
    parser = argparse.ArgumentParser(description="调用 qvq 做组内图文一对一匹配，最终输出扁平 {source, caption} 列表")
    parser.add_argument("--groups", required=True, help="包含多组对象的 JSON（每组含 sources 与 union）")
    parser.add_argument("--output", required=True, help="输出扁平 JSON 文件路径")
    parser.add_argument("--image-root", default="xfdata/converted_pngs", help="图片根目录（默认 xf_data/converted_pngs）")
    parser.add_argument("--model", default="qvq-max-latest", help="qvq 模型名（默认 qvq-max）")
    parser.add_argument("--base-url", default="https://dashscope.aliyuncs.com/compatible-mode/v1", help="OpenAI 兼容 API base_url")
    parser.add_argument("--api-key-env", default="sk-8863180d60534ea7909fb1745a27fac8", help="保存 API Key 的环境变量名")
    parser.add_argument("--limit-groups", type=int, default=0, help="只处理前 N 组（调试用）")
    parser.add_argument("--require-all-images-exist", action="store_true",
                        help="开启后：若某组存在缺图则跳过该组；默认：遇缺图直接报错退出")
    args = parser.parse_args()

    api_key = args.api_key_env
    if not api_key:
        print(f"环境变量 {args.api_key_env} 未设置。", file=sys.stderr)
        sys.exit(1)

    client = OpenAI(api_key=api_key, base_url=args.base_url)

    with open(args.groups, "r", encoding="utf-8") as f:
        groups = json.load(f)
    if not isinstance(groups, list):
        print("输入 JSON 顶层必须是列表。", file=sys.stderr)
        sys.exit(2)

    if args.limit_groups and args.limit_groups > 0:
        groups = groups[:args.limit_groups]

    flat_results = []  # 只存 {source, caption}

    count = 0
    for gi, group in enumerate(groups):
        print(count)
        count = count+1
        sources = group.get("sources") or group.get("Sources") or []
        candidates = group.get("union") or []
        if not sources:
            continue

        data_urls = []
        index_to_source = []
        missing = []

        for src in sources:
            png_path = map_source_to_png_path(src, args.image_root)
            if not os.path.exists(png_path):
                missing.append(png_path)
                continue
            data_urls.append(to_data_url_png(png_path))
            index_to_source.append(src)

        if missing:
            msg = f"[组 {gi}] 缺少图片文件：\n" + "\n".join(missing)
            if args.require_all_images_exist:
                print(msg + "\n该组已跳过。", file=sys.stderr)
                continue
            else:
                print(msg, file=sys.stderr)
                sys.exit(3)

        if not candidates:
            # 没有候选：全部置空
            for s in index_to_source:
                flat_results.append({"source": s, "caption": ""})
            print(f"[组 {gi}] 无候选，已输出空 caption。")
            continue

        messages = build_messages(data_urls, candidates)

        try:
            raw = call_qvq(client, args.model, messages)
        except Exception as e:
            print(f"[组 {gi}] qvq 调用失败：{e}", file=sys.stderr)
            sys.exit(4)

        try:
            parsed = safe_json_loads(raw)
        except Exception as e:
            print(f"[组 {gi}] 解析失败：{e}\n返回：\n{raw}", file=sys.stderr)
            sys.exit(5)

        assignments = parsed.get("assignments")
        if not isinstance(assignments, list):
            print(f"[组 {gi}] 返回缺少 'assignments' 数组。\n返回：\n{raw}", file=sys.stderr)
            sys.exit(6)

        # 校验与落地
        candidates_set = set(candidates)
        used = set()
        got_idx = set()

        for item in assignments:
            if not isinstance(item, dict):
                continue
            if "index" not in item:
                continue
            idx = item["index"]
            if not isinstance(idx, int) or idx < 0 or idx >= len(index_to_source):
                continue
            cap = item.get("caption", "")
            got_idx.add(idx)

            if cap and cap not in candidates_set:
                cap = ""
            if cap and cap in used:
                cap = ""
            if cap:
                used.add(cap)

            flat_results.append({
                "source": index_to_source[idx],
                "caption": cap
            })

        # 补齐未出现的 index
        for idx in range(len(index_to_source)):
            if idx not in got_idx:
                flat_results.append({
                    "source": index_to_source[idx],
                    "caption": ""
                })

        print(f"[组 {gi}] 完成：{len(index_to_source)} 条。")

    with open(args.output, "w", encoding="utf-8") as f:
        json.dump(flat_results, f, ensure_ascii=False, indent=2)

    print(f"全部完成：共写入 {len(flat_results)} 条 -> {args.output}")

if __name__ == "__main__":
    main()
